package com.gin.nebula

import com.facebook.thrift.protocol.TCompactProtocol
import com.gin.utils.ResultSetUtil
import com.vesoft.nebula.client.graph.data.ResultSet
import com.vesoft.nebula.connector.connector.{NebulaDataFrameReader, NebulaDataFrameWriter}
import com.vesoft.nebula.connector.nebula.GraphProvider
import com.vesoft.nebula.connector.{Address, NebulaConnectionConfig, ReadNebulaConfig, WriteNebulaVertexConfig}
import org.apache.log4j.BasicConfigurator
import org.apache.spark.SparkConf
import org.apache.spark.sql.{Encoders, Row, SparkSession}

object NebulaReadWriteDemo {

  BasicConfigurator.configure()
  val nebula_ip: String = "192.168.1.102"
  val queryNGql: String =
    """
use xwqy_element_shareholding_graph;
GET SUBGRAPH WITH PROP 10 STEPS
FROM VID_REPLACE_PLACEHOLDER IN e_invest
YIELD VERTICES AS nodes, EDGES AS relationships;
"""

  def main(args: Array[String]): Unit = {



    val sparkConf = new SparkConf
    sparkConf
      .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
      .registerKryoClasses(Array[Class[_]](classOf[TCompactProtocol]))
    val sparkSession: SparkSession = SparkSession
      .builder()
      .master("local")
      .config(sparkConf)
      .getOrCreate()

    //仅打印错误日志
    sparkSession.sparkContext.setLogLevel("ERROR")

    //读取企业
    val readConfigs = getNebulaReadNebulaConfig()
    val vertex = sparkSession.read.nebula(readConfigs._1, readConfigs._2).loadVerticesToDF()
    vertex.printSchema()
    vertex.show()

    //加工
    val entTidAdd = vertex.mapPartitions(
      (partitionIter) => {
        new Iterator[(String, java.lang.Long, String, String)] {
          //创建nebula连接
          println(s"--------nebula-connect--------")
          val graphProvider = getGraphProvider()
          graphProvider.switchSpace("root", "nebula", "xwqy_element_shareholding_graph")

          //返回是否有下一个元素,并在没有元素时释放nebula连接
          override def hasNext: Boolean = if (!partitionIter.hasNext) {
            try {
              println("--------close-nebula-connect--------")
              graphProvider.releaseGraphClient()
              graphProvider.close()
            } catch {
              case e: Exception =>
                println("--------close-nebula-connect-error--------", e)
            }
            false
          } else {
            true
          }

          //加工 查询图数据库
          override def next(): (String, java.lang.Long, String, String) = {
            val row: Row = partitionIter.next()

            val vid = row.getAs[String]("_vertexId")
            val tid = row.getAs[Long]("tid")

            //查询图数据库操作
            val vertexResultSet: ResultSet =
              graphProvider.submit(queryNGql.replace("VID_REPLACE_PLACEHOLDER", vid.toString))
            println(s"--------nebula-select-$vertexResultSet--------")
            ResultSetUtil.parseResultToGraph(vertexResultSet, vid.toLong)

            (vid, tid,
              row.getAs[String]("tname"), row.getAs[String]("uniscid"))
          }

        }
      }
    )(Encoders.tuple(Encoders.STRING, Encoders.LONG, Encoders.STRING, Encoders.STRING))

    entTidAdd.foreach(row => {
      println(row)
    })

    //重新写入企业
    val writeVertexConfig = getWriteNebulaVertexConfig()
    val df = entTidAdd.toDF("_vertexId", "tid", "tname", "uniscid")
    df.show()
    df.write.nebula(writeVertexConfig._1, writeVertexConfig._2).writeVertices()

    println("nebula tag end...")
    sparkSession.stop()

  }

  def getGraphProvider(): GraphProvider = {
    val addresses: List[Address] = List(new Address(nebula_ip, 9669))
    new GraphProvider(addresses, 3000)
  }

  def getNebulaReadNebulaConfig(): (NebulaConnectionConfig, ReadNebulaConfig) = {
    val config =
      NebulaConnectionConfig
        .builder()
        .withMetaAddress(nebula_ip+":9559")
        .withConenctionRetry(2)
        .build()
    val nebulaReadVertexConfig: ReadNebulaConfig = ReadNebulaConfig
      .builder()
      .withSpace("xwqy_element_shareholding_graph")
      .withLabel("t_r_ent")
      .withNoColumn(false)
      .withReturnCols(List("tid", "tname", "uniscid"))
      .withLimit(10)
      .withPartitionNum(2)
      .build()
    (config, nebulaReadVertexConfig)
  }

  def getWriteNebulaVertexConfig(): (NebulaConnectionConfig, WriteNebulaVertexConfig) = {
    // connection config without ssl
    val config =
      NebulaConnectionConfig
        .builder()
        .withMetaAddress(nebula_ip+":9559")
        .withGraphAddress(nebula_ip+":9669")
        .withConenctionRetry(2)
        .build()

    val nebulaWriteVertexConfig: WriteNebulaVertexConfig = WriteNebulaVertexConfig
      .builder()
      .withSpace("xwqy_element_shareholding_graph")
      .withTag("t_r_ent")
      .withVidField("_vertexId")
      .withVidAsProp(false)
      .withBatch(1000)
      .build()

    (config, nebulaWriteVertexConfig)
  }
}
